package com.huawei.deveco.common.security.encrypt;

import com.huawei.deveco.common.security.util.*;

import java.io.File;
import java.io.IOException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;

public class KeyManager {
    private static final int XOR_OP_COUNT = 2;

    private static final int COMPONENT_COUNT = 3;

    private static final int KEY_LENGTH = 16;

    private static final String[] COMPONENT_DIR_NAMES = new String[]{"fd", "ac", "ce"};

    private static final int ITERATE_COUNT = 10000;

    private static final byte[] COMPONENT = new byte[]{
            49, -13, 9, 115, -42, -81, 91, -72, -45, -66,
            -79, 88, 101, -125, -64, 119};

    private static final String ROOT_FOLDER = "material";

    private static final String DS_STORE_FILE = ".DS_Store";

    private final CryptoKit cryptoKit = CryptoKitFactory.getInstance("AES-GCM-128");

    public byte[] readOrCreateKey(String materialDir) {
        if (StringUtil.isEmpty(materialDir))
            throw new IllegalArgumentException("MaterialDir should not be empty");
        String rootMaterial = materialDir + File.separator + "material";
        File rootDir = checkOrCreateDirIfNecessary(rootMaterial);
        File[] componentFiles = rootDir.listFiles();
        if (componentFiles == null || componentFiles.length == 0)
            return generateKey(rootDir);
        checkComponent(componentFiles);
        return readKey(rootDir);
    }

    private void checkComponent(File[] componentFiles) {
        Set<String> fileNames = (Set<String>) Stream.<File>of(componentFiles).map(file -> file.getName()).collect(Collectors.toSet());
        for (String componentName : COMPONENT_DIR_NAMES) {
            if (!fileNames.contains(componentName))
                throw new IllegalStateException(
                        String.format(Locale.ENGLISH, "Component in the [%s] dir has been damaged!", new Object[]{"material"}));
        }
    }

    private File checkOrCreateDirIfNecessary(String materialDir) {
        File rootDir = new File(materialDir);
        if (rootDir.isDirectory())
            return rootDir;
        if (!rootDir.exists())
            try {
                FilePermissionUtil.createMaterialDirsWithPermission(rootDir.toPath());
                return rootDir;
            } catch (IOException e) {
                throw new IllegalStateException("Create [material] dir failed.");
            }
        throw new IllegalStateException("File [material] should be a directory!");
    }

    private byte[] generateKey(File rootDir) {
        SecureRandom secureRandom = DRBG.getSecureRandom();
        List<byte[]> components = generateComponents(secureRandom);
        byte[] salt = generateBytes(secureRandom);
        try {
            byte[] rootKey = deriveRootKey(components, salt);
            byte[] workKey = generateBytes(secureRandom);
            byte[] workKeyMaterial = this.cryptoKit.encrypt(rootKey, workKey);
            saveComponents(rootDir, components);
            saveSalt(rootDir, salt);
            saveWorkMaterial(rootDir, workKeyMaterial);
            return workKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException |
                 com.huawei.deveco.common.security.exception.CryptoException | IOException e) {
            throw new IllegalStateException("generate material failed");
        }
    }

    private void saveComponents(File rootDir, List<byte[]> components) throws IOException {
        for (int i = 0; i < components.size(); i++)
            saveBytes(assembleFile(new File(rootDir, COMPONENT_DIR_NAMES[0] + File.separator + i)), components.get(i));
    }

    private void saveSalt(File rootDir, byte[] salt) throws IOException {
        saveBytes(assembleFile(new File(rootDir, COMPONENT_DIR_NAMES[1])), salt);
    }

    private void saveWorkMaterial(File rootDir, byte[] workKeyMaterial) throws IOException {
        saveBytes(assembleFile(new File(rootDir, COMPONENT_DIR_NAMES[2])), workKeyMaterial);
    }

    private File assembleFile(File file) {
        return new File(file, generateFileName());
    }

    private String generateFileName() {
        return UUID.randomUUID().toString().replace("-", "");
    }

    private void saveBytes(File componentDir, byte[] bytes) throws IOException {
        FilePermissionUtil.createMaterialFileOrWriteBytesWithPermission(componentDir.toPath(), bytes);
    }

    private List<byte[]> getFullComponents(List<byte[]> components) {
        List<byte[]> fullComponents = (List) new ArrayList<>(components);
        fullComponents.add(COMPONENT);
        return fullComponents;
    }

    private byte[] deriveRootKey(List<byte[]> components, byte[] salt) throws NoSuchAlgorithmException, InvalidKeySpecException {
        List<byte[]> fullComponents = getFullComponents(components);
        char[] finalComponent = ByteUtil.bytes2chars(xorComponents(fullComponents));
        PBEKeySpec spec = new PBEKeySpec(finalComponent, salt, 10000, 128);
        SecretKeyFactory skf = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
        return skf.generateSecret(spec).getEncoded();
    }

    private byte[] xorComponents(List<byte[]> components) {
        checkComponents(components);
        byte[] result = ByteUtil.xor(components.get(0), components.get(1));
        int componentsSize = components.size();
        for (int i = 2; i < componentsSize; i++)
            result = ByteUtil.xor(result, components.get(i));
        return result;
    }

    private void checkComponents(List<byte[]> components) {
        int componentsSize = components.size();
        if (componentsSize < 2)
            throw new IllegalStateException("Components needs at least 2");
        for (byte[] bytes : components) {
            if (bytes.length != 16)
                throw new IllegalStateException("Components length is not same!");
        }
    }

    private List<byte[]> generateComponents(SecureRandom secureRandom) {
        List<byte[]> components = (List) new ArrayList<>();
        for (int i = 0; i < 3; i++)
            components.add(generateBytes(secureRandom));
        return components;
    }

    private byte[] generateBytes(SecureRandom secureRandom) {
        byte[] bytes = new byte[16];
        secureRandom.nextBytes(bytes);
        return bytes;
    }

    private byte[] readKey(File rootDir) {
        try {
            List<byte[]> components = readComponents(rootDir);
            byte[] salt = readSalt(rootDir);
            byte[] rootKey = deriveRootKey(components, salt);
            byte[] workMaterial = readWorkMaterial(rootDir);
            return this.cryptoKit.decrypt(rootKey, workMaterial);
        } catch (IOException | NoSuchAlgorithmException | InvalidKeySpecException |
                 com.huawei.deveco.common.security.exception.CryptoException e) {
            throw new IllegalStateException("Read Material failed!");
        }
    }

    private List<byte[]> readComponents(File rootDir) throws IOException {
        File componentRootDir = new File(rootDir, COMPONENT_DIR_NAMES[0]);
        if (!componentRootDir.exists())
            throw new IllegalStateException("Material not found!");
        File[] componentDirs = componentRootDir.listFiles();
        if (componentDirs == null)
            throw new IllegalStateException("Components state is illegal!");
        if (OsTypeUtil.isMac())
            componentDirs = (File[]) Arrays.<File>stream(componentDirs).filter(file -> !".DS_Store".equals(file.getName())).toArray(x$0 -> new File[x$0]);
        if (componentDirs.length != 3)
            throw new IllegalStateException("Components state is illegal!");
        List<byte[]> components = (List) new ArrayList<>();
        for (File dir : componentDirs)
            components.add(readBytes(dir));
        return components;
    }

    private byte[] readSalt(File rootDir) throws IOException {
        return readBytes(new File(rootDir, COMPONENT_DIR_NAMES[1]));
    }

    private byte[] readWorkMaterial(File rootDir) throws IOException {
        return readBytes(new File(rootDir, COMPONENT_DIR_NAMES[2]));
    }

    private byte[] readBytes(File componentDir) throws IOException {
        File[] component = componentDir.listFiles();
        if (component == null || component.length != 1)
            throw new IllegalStateException("Sub component state is illegal!");
        return FileUtil.readFileBytes(component[0]);
    }
}
